import os
from triton.backends.compiler import GPUTarget

blackwell = GPUTarget("cuda", 100, 32)
hopper = GPUTarget("cuda", 90, 32)
ampere = GPUTarget("cuda", 80, 32)
hip_rdna3 = GPUTarget("hip", "gfx1100", 32)
hip_rdna4 = GPUTarget("hip", "gfx1200", 32)
hip_cdna3 = GPUTarget("hip", "gfx942", 64)
hip_cdna4 = GPUTarget("hip", "gfx950", 64)

gpus = {
    "BLACKWELL": blackwell,
    "HOPPER": hopper,
    "AMPERE": ampere,
    "HIP_RDNA3": hip_rdna3,
    "HIP_RDNA4": hip_rdna4,
    "HIP_CDNA3": hip_cdna3,
    "HIP_CDNA4": hip_cdna4,
}

try:
    GPU = os.environ['GPU'].upper()
    if GPU in gpus:
        TARGET = gpus[GPU]
        os.environ['GPU_NAME'] = GPU
    else:
        assert False, f"unsupported {GPU}"
except KeyError:
    TARGET = blackwell
    os.environ['GPU_NAME'] = "blockwell"
